import cv2 as cv
import numpy as np
import os
from PIL import Image
from glob import glob
import matplotlib.pyplot as plt


def video2frames(video_path):
    frames = []
    video = cv.VideoCapture(video_path)
    frame_num = video.get(7)
    print('video: {} has {:d} frames\n'.format(video_path, int(frame_num)))
    ret, frame = video.read()

    while ret and video.isOpened():
        frames.append(frame)
        ret, frame = video.read()

    print('done\n')

    return frames


def cal_for_frames(frames):
    flow = []
    # prev = cv.imread(frames[0])
    prev = frames[0]
    prev = cv.cvtColor(prev, cv.COLOR_BGR2GRAY)

    for i, frame_curr in enumerate(frames[1:]):
        curr = frame_curr
        curr = cv.cvtColor(curr, cv.COLOR_BGR2GRAY)

        tmp_flow = compute_TVL1(prev, curr)
        flow.append(tmp_flow)
        prev = curr

    return flow


def compute_TVL1(prev, curr, bound=15):
    TVL1 = cv.optflow.DualTVL1OpticalFlow_create()
    flow = TVL1.calc(prev, curr, None)
    assert flow.dtype == np.float32

    flow = (flow + bound) * (255.0 / (2 * bound))
    flow = np.round(flow).astype(int)
    np.clip(flow, 0, 255, out=flow)

    return flow


def save_flow(video_flows, flow_path):
    if not os.path.exists(os.path.join(flow_path, 'u')):
        os.mkdir(os.path.join(flow_path, 'u'))

    if not os.path.exists(os.path.join(flow_path, 'v')):
        os.mkdir(os.path.join(flow_path, 'v'))

    for i, flow in enumerate(video_flows):
        cv.imwrite(os.path.join(flow_path, 'u', 'u_{:05d}.jpg'.format(i)),
                   flow[:, :, 0])
        cv.imwrite(os.path.join(flow_path, 'v', 'v_{:05d}.jpg'.format(i)),
                   flow[:, :, 1])


def extract_flow(video_path, flow_path):
    frames = video2frames(video_path)
    flow = cal_for_frames(frames)
    save_flow(flow, flow_path)
    print('Complete: {}\n'.format(flow_path))
    return True


def display(path):
    plt.ion()

    files = os.listdir(path)
    files.sort()
    for file in files:
        img = Image.open(os.path.join(path, file)).convert('P')
        plt.title(file)
        plt.imshow(img)
        plt.pause(0.1)

    plt.ioff()
    plt.show()
    return


def display(paths, frames):
    plt.ion()

    fig, axs = plt.subplots(1, 2)

    file1 = os.listdir(paths[0])
    file2 = os.listdir(paths[1])

    file1.sort()
    file2.sort()

    def animation(ax, path, file):
        path = os.path.join(path, file)
        img = Image.open(path).convert('P')
        ax.set_title(file)
        ax.set_xticks([])
        ax.set_yticks([])
        ax.imshow(img)
        return

    for i in range(len(file1)):
        animation(axs[0], paths[0], file1[i])
        animation(axs[1], paths[1], file2[i])
        plt.pause(0.1)

    plt.ioff()
    plt.show()


def clean_jpg(path):
    imgs = glob(os.path.join(path, '*.jpg'))
    for img in imgs:
        os.remove(img)

    print('clean all images in {}\n'.format(path))


if __name__ == '__main__':
    video_path = "C:/Users/18873/Desktop/videos/v_ApplyEyeMakeup_g01_c01.avi"
    save_path = '../assets/'
    extract_flow(video_path, save_path)
    # display(os.path.join(save_path, 'u'))
    # clean_jpg(os.path.join(save_path, 'v'))
    # clean_jpg(os.path.join(save_path, 'u'))
    paths = [os.path.join(save_path, 'u'), os.path.join(save_path, 'v')]
    # display(paths, frames=None)
